#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Python version: 3.6

import torch
from torch import nn
import torch.nn.functional as F
from efficientnet_pytorch import EfficientNet



class EfficientNetModel(nn.Module):
    def __init__(self, args):
        super(EfficientNetModel, self).__init__()
        self.efficientnet = EfficientNet.from_pretrained('efficientnet-b0')
        self.num_classes = args.num_classes
        self.fc = nn.Linear(1280, self.num_classes)

    def forward(self, x):
        x = self.efficientnet.extract_features(x)
        x = F.avg_pool2d(x, x.shape[2:])
        x = x.view(x.shape[0], -1)
        x = self.fc(x)
        return x
